(*
 * Functions for indexing and iterating over the states of sets of variables
 *)

(* State index --> state of variables *)
let get_offsets ranges = 
  let offsets = Array.create (Array.length ranges) 0 in
  offsets.(Array.length offsets - 1) <- 1 ;
  for i = Array.length offsets - 2 downto 0 do
    offsets.(i) <- ranges.(i+1) * offsets.(i+1)
  done ; offsets

let idx_to_varstate ranges state_idx =
  let offsets = get_offsets ranges in
  let state = Array.create (Array.length ranges) (-1) in
  let remainder = ref state_idx in
  for i = 0 to Array.length ranges - 1 do 
    state.(i) <- !remainder / offsets.(i) ;
    remainder := !remainder mod offsets.(i)
  done ;
  state

(* State of variables --> state index *)
let varstate_to_idx ranges values =
  let idx = ref 0 in
  for i = 0 to Array.length ranges - 1 do 
    idx := !idx * ranges.(i) + values.(i)
  done ;
  !idx

(* State of *all* variables --> state index *)
let state_to_idx vars ranges values =
  let idx = ref 0 in
  for i = 0 to Array.length ranges - 1 do 
    idx := !idx * ranges.(i) + values.(vars.(i))
  done ;
  !idx


exception NoMoreStates

(* Helper function for iterating over all states of the parent variables *)
(* TODO -- this interface is a bit awkward. *)
let rec incstate schema varstate parents =
  match parents with
    x :: l ->
      if varstate.(x) < schema.(x) - 1 then
        varstate.(x) <- varstate.(x) + 1
      else
       (varstate.(x) <- 0;
        incstate schema varstate l)
  | [] -> raise NoMoreStates


let map_state schema vars f =
  let varstate = Array.create (Array.length schema) (-1) in
  List.iter (fun i -> varstate.(i) <- 0) vars ;
  let results = ref [] in
  (try 
    while true do
      results := (f varstate) :: !results;
      incstate schema varstate vars
    done
  with NoMoreStates -> ()) ;
  !results

let iter_state schema vars f = ignore (map_state schema vars f)
